import numpy as np
import matplotlib.pyplot as plt

class LQRController:
    def __init__(self, A, B, Q, R):
        self.A = A
        self.B = B
        self.Q = Q
        self.R = R
        self.measurement_history = []
        self.setpoint_history = []

    def fuzzy_adjustment(self, error, error_dot):
        q_factor, r_factor = self.calculate_fuzzy_factors(error, error_dot)
        adjusted_Q = self.Q * q_factor
        adjusted_R = self.R * r_factor
        return adjusted_Q, adjusted_R

    def calculate_fuzzy_factors(self, error, error_dot):
        e_membership = fuzzy_membership(error, [-1, 0, 0.5, 1, 2])
        e_dot_membership = fuzzy_membership(error_dot, [-0.5, 0, 0.25, 0.5, 1])

        # 根据模糊集合定义的规则，返回不同的q_factor和r_factor
        if e_membership[0] and e_dot_membership[0]:  # 非常低误差和非常低误差变化速率
            q_factor = 0.5
            r_factor = 0.5
        elif e_membership[0] and e_dot_membership[1]:  # 非常低误差和低误差变化速率
            q_factor = 0.7
            r_factor = 0.6
        elif e_membership[1] and e_dot_membership[0]:  # 低误差和非常低误差变化速率
            q_factor = 0.8
            r_factor = 0.7
        elif e_membership[1] and e_dot_membership[1]:  # 低误差和低误差变化速率
            q_factor = 1.0
            r_factor = 0.8
        elif e_membership[2] and e_dot_membership[2]:  # 中等误差和中等误差变化速率
            q_factor = 1.2
            r_factor = 1.0
        elif e_membership[3] and e_dot_membership[3]:  # 高误差和高误差变化速率
            q_factor = 1.5
            r_factor = 1.3
        elif e_membership[4] and e_dot_membership[4]:  # 非常高误差和非常高误差变化速率
            q_factor = 2.0
            r_factor = 1.8
        elif e_membership[3] and e_dot_membership[4]:  # 高误差和非常高误差变化速率
            q_factor = 1.7
            r_factor = 1.5
        elif e_membership[4] and e_dot_membership[3]:  # 非常高误差和高误差变化速率
            q_factor = 1.9
            r_factor = 1.6
        else:  # 其他组合
            q_factor = 1.0
            r_factor = 1.0

        return q_factor, r_factor

    def lqr_solve(self, Q=None, R=None):
        Q = Q if Q is not None else self.Q
        R = R if R is not None else self.R

        if np.isscalar(R):
            R = np.array([[R]])
        elif R.ndim == 1:
            R = np.diag(R)

        n = self.A.shape[0]
        m = self.B.shape[1]

        Z = np.vstack([
            np.hstack([self.A, -self.B @ np.linalg.inv(R) @ self.B.T]),
            np.hstack([-Q, -self.A.T])
        ])
        eigvals, eigvecs = np.linalg.eig(Z)

        stable_indices = np.where(np.real(eigvals) < 0)[0]
        eigvecs_stable = eigvecs[:, stable_indices]

        X = eigvecs_stable[:n, :]
        Y = eigvecs_stable[n:, :]

        P = np.real(Y @ np.linalg.inv(X))

        K = np.linalg.inv(R) @ self.B.T @ P

        eigVals, _ = np.linalg.eig(self.A - self.B @ K)

        return K, P, eigVals

    def update(self, setpoint, measurement):
        self.measurement_history.append(measurement)
        self.setpoint_history.append(setpoint)
        error = setpoint - measurement
        error_dot = np.zeros_like(error)
        adjusted_Q, adjusted_R = self.fuzzy_adjustment(error, error_dot)
        K, _, _ = self.lqr_solve(adjusted_Q, adjusted_R)
        u = K @ error
        return u

    def plot_error(self, title='none'):
        self.setpoint_history = np.array(self.setpoint_history)
        self.measurement_history = np.array(self.measurement_history)
        for i in range(len(self.setpoint_history[0])):
            fig, ax = plt.subplots()
            ax.plot([i / 100 for i in range(len(self.setpoint_history[:, i]))], self.setpoint_history[:, i],
                    label='setpoint')
            ax.plot([i / 100 for i in range(len(self.measurement_history[:, i]))], self.measurement_history[:, i],
                    label='measurement')
            ax.set_xlabel('time')
            ax.set_ylabel('value')
            ax.set_title(title)
            ax.legend()

    def plot_fuzzy_rules(self):
        # Plot for error
        error_range = [-1, 0, 0.5, 1, 2]  # Range for error
        plt.figure(figsize=(12, 5))

        x = np.linspace(error_range[0] - 1, error_range[4] + 1, 100)
        for i, membership_func in enumerate(['Very Low', 'Low', 'Medium', 'High', 'Very High']):
            if membership_func == 'Very Low':
                y = np.where(x < error_range[0], 1, np.where(x < error_range[1], (error_range[1] - x) / (error_range[1] - error_range[0]), 0))
            elif membership_func == 'Low':
                y = np.where(x < error_range[0], 0, np.where(x < error_range[1], (x - error_range[0]) / (error_range[1] - error_range[0]), np.where(x < error_range[2], (error_range[2] - x) / (error_range[2] - error_range[1]), 0)))
            elif membership_func == 'Medium':
                y = np.where(x < error_range[1], 0, np.where(x < error_range[2], (x - error_range[1]) / (error_range[2] - error_range[1]), np.where(x < error_range[3], (error_range[3] - x) / (error_range[3] - error_range[2]), 0)))
            elif membership_func == 'High':
                y = np.where(x < error_range[2], 0, np.where(x < error_range[3], (x - error_range[2]) / (error_range[3] - error_range[2]), np.where(x < error_range[4], (error_range[4] - x) / (error_range[4] - error_range[3]), 0)))
            elif membership_func == 'Very High':
                y = np.where(x < error_range[3], 0, np.where(x < error_range[4], (x - error_range[3]) / (error_range[4] - error_range[3]), 1))
            plt.plot(x, y, label=f'{membership_func} Membership - Error')

        plt.title('Error Membership Functions')
        plt.xlabel('Error Value')
        plt.ylabel('Membership Degree')
        plt.legend()
        plt.grid(True)

        # Plot for error_dot
        error_dot_range = [-0.5, 0, 0.25, 0.5, 1]  # Range for error_dot
        plt.figure(figsize=(12, 5))

        x = np.linspace(error_dot_range[0] - 0.1, error_dot_range[4] + 0.1, 100)
        for i, membership_func in enumerate(['Very Low', 'Low', 'Medium', 'High', 'Very High']):
            if membership_func == 'Very Low':
                y = np.where(x < error_dot_range[0], 1, np.where(x < error_dot_range[1], (error_dot_range[1] - x) / (error_dot_range[1] - error_dot_range[0]), 0))
            elif membership_func == 'Low':
                y = np.where(x < error_dot_range[0], 0, np.where(x < error_dot_range[1], (x - error_dot_range[0]) / (error_dot_range[1] - error_dot_range[0]), np.where(x < error_dot_range[2], (error_dot_range[2] - x) / (error_dot_range[2] - error_dot_range[1]), 0)))
            elif membership_func == 'Medium':
                y = np.where(x < error_dot_range[1], 0, np.where(x < error_dot_range[2], (x - error_dot_range[1]) / (error_dot_range[2] - error_dot_range[1]), np.where(x < error_dot_range[3], (error_dot_range[3] - x) / (error_dot_range[3] - error_dot_range[2]), 0)))
            elif membership_func == 'High':
                y = np.where(x < error_dot_range[2], 0, np.where(x < error_dot_range[3], (x - error_dot_range[2]) / (error_dot_range[3] - error_dot_range[2]), np.where(x < error_dot_range[4], (error_dot_range[4] - x) / (error_dot_range[4] - error_dot_range[3]), 0)))
            elif membership_func == 'Very High':
                y = np.where(x < error_dot_range[3], 0, np.where(x < error_dot_range[4], (x - error_dot_range[3]) / (error_dot_range[4] - error_dot_range[3]), 1))
            plt.plot(x, y, label=f'{membership_func} Membership - Error_dot')

        plt.title('Error_dot Membership Functions')
        plt.xlabel('Error_dot Value')
        plt.ylabel('Membership Degree')
        plt.legend()
        plt.grid(True)


def fuzzy_membership(value, range):
    very_low, low, medium, high, very_high = range
    membership = np.zeros(5)

    if isinstance(value, np.ndarray):
        for i, v in enumerate(value):
            if v < very_low:
                membership[i] = 1
            elif very_low <= v < low:
                membership[i] = (low - v) / (low - very_low)
            elif low <= v < medium:
                membership[i] = (medium - v) / (medium - low)
            elif medium <= v < high:
                membership[i] = (high - v) / (high - medium)
            elif high <= v < very_high:
                membership[i] = (very_high - v) / (very_high - high)
            else:
                membership[i] = 0
    else:
        if value < very_low:
            membership = [1, 0, 0, 0, 0]
        elif very_low <= value < low:
            membership = [(low - value) / (low - very_low), 1 - abs(value - low) / (low - very_low), 0, 0, 0]
        elif low <= value < medium:
            membership = [0, (medium - value) / (medium - low), 1 - abs(value - medium) / (medium - low), 0, 0]
        elif medium <= value < high:
            membership = [0, 0, (high - value) / (high - medium), 1 - abs(value - high) / (high - medium), 0]
        elif high <= value < very_high:
            membership = [0, 0, 0, (very_high - value) / (very_high - high), 1 - abs(value - very_high) / (very_high - high)]
        else:
            membership = [0, 0, 0, 0, 1]

    return membership


if __name__ == '__main__':
    A = np.array([[0., 1.], [-1., -1.]])
    B = np.array([[0.], [1.]])
    Q = np.array([[1., 0.], [0., 1.]])
    R = np.array([[1.]])

    pos_controller = LQRController(A, B, Q, R)
    setpoint_pos = np.array([1, 0])
    measurement_pos = np.array([0, 0])

    vel_setpoint = pos_controller.update(setpoint_pos, measurement_pos)
    print("Velocity setpoint:", vel_setpoint)

    pos_controller.plot_error("Position Control Error")
    pos_controller.plot_fuzzy_rules()
    plt.show()
